package dnavsim;
import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.BufferedWriter;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.FileWriter;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Properties;

import org.apache.log4j.Level;
import org.apache.log4j.Logger;

import dnavsim.bed.BEDFeature;
import dnavsim.variant.Cnv;
import dnavsim.variant.Deletion;
import dnavsim.variant.Insertion;
import dnavsim.variant.Snp;
import dnavsim.variant.Variant;
import dnavsim.variant.VariantHandler;

public class DNAVSim {

	private Logger log;

	public static final String[] CHROMOSOMES = new String[] { "1", "2", "3",
			"4", "5", "6", "7", "8", "9", "10", "11", "12", "13", "14", "15",
			"16", "17", "18", "19", "20", "21", "22", "23", "X", "Y", "MT" };
	
	public static String PROGRAM_NAME = "DNA Variant Simulator";
	public static String PROGRAM_VERSION = "0.2 beta";
	public static String CONFIG_FILENAME = "config/dnavsim.cfg";

	/**
	 * String used to tag temporary files that should be deleted 
	 * at the end of the process. 
	 */
	public static String TMP_FILNAME_TAG = "__tmp__";
	
	public static String FASTA_FILE_COLUMNS_KEY = "out.fasta.columns";
	public static String VARIANT_SNP_FREQ_KEY = "variant.snp.freq";
	public static String VARIANT_INS_FREQ_KEY = "variant.ins.freq";
	public static String VARIANT_DEL_FREQ_KEY = "variant.del.freq";
	public static String VARIANT_CNV_III_FREQ_KEY = "variant.cnv_III.freq";
	public static String INDEL_LENGTH_MIN_KEY = "indel.length.min";
	public static String INDEL_LENGTH_MAX_KEY = "indel.length.max";
	public static String REGION_SET_COUNT_KEY = "regionset.count";
	
	public static String REGION_SET_NAME_KEY = "regionset.#.name"; 
	public static String REGION_SET_VARIANT_FREQ_KEY = "regionset.#.variantFreq"; 
	public static String REGION_SET_BED_KEY = "regionset.#.bed"; 
	
	/**
	 * Region set reserved names
	 */
	public static String REGION_SET_NAME_GENE = "gene";
	public static String REGION_SET_NAME_INTRON = "intron";
	public static String REGION_SET_NAME_EXON = "exon";
	
	public static String REGION_SET_DEFAULT_FREQ_KEY = "regionset.default.variant.freq"; 
	
	public static String CONFIG_KEYS[] = { VARIANT_SNP_FREQ_KEY,
			VARIANT_INS_FREQ_KEY, VARIANT_DEL_FREQ_KEY,
			VARIANT_CNV_III_FREQ_KEY, INDEL_LENGTH_MIN_KEY,
			INDEL_LENGTH_MAX_KEY, FASTA_FILE_COLUMNS_KEY, REGION_SET_COUNT_KEY,
			REGION_SET_DEFAULT_FREQ_KEY };

	/**
	 * Configuration parameters loaded from configuration file
	 */
	private int FASTA_FILE_COLUMNS;
	private double SNP_FREQ;
	private double INS_FREQ;
	private double DEL_FREQ;
	private double CNV_III_FREQ;
	private double DEFAULT_REGION_FREQ;	
	private int INDEL_MIN_LENGTH;
	private int INDEL_MAX_LENGTH;
	private RegionSet[] REGION_SETS;
	private RegionSet GENES;
	private RegionSet INTRONS;
	private RegionSet EXONS;
	
	// The configuration file 
	private Properties cfg;
	
	/**
	 * Parse a FASTA header line to find out the chromosome name
	 * corresponding to the subsequent sequence
	 * 
	 * @param line the FASTA header line (should start by ">"
	 * @return the chromosome name (e.g. 1, 2, 3, ..., X, or Y)
	 */
	private String parseFastaHeaderLine(String line) {
		
		int index = -1;
		String result = null;
		if ((index = line.indexOf("chr")) != -1) {
			result = line.substring(index+3);
		}
		else 
		if ((index = line.indexOf("chrom")) != -1) {
			result = line.substring(index+5);
		}
		else 
			result = line.substring(1).trim();
		
		// Remove optional loci 
		if (result.indexOf(':') != -1)
			result = result.substring(0, result.indexOf(':'));
		
		return result;
	}

	/**
	 * load the configuration file and do a sanity check of the parameters
	 */
	private void loadConfiguration() {

		log.info("Loading properties from 'dnavsim.cfg' file...");
		try {
			cfg = new Properties();
			cfg.load(new FileInputStream(CONFIG_FILENAME));
		} catch (FileNotFoundException e1) {
			log.fatal("An error occured (File not found exception) while loading configuration file '" + CONFIG_FILENAME + "'! Exiting...");
			System.exit(1);
		} catch (IOException e1) {
			log.fatal("An IO error occured while loading configuration file '" + CONFIG_FILENAME + "'! Exiting...");
			System.exit(1);
		}
		for (int i=0; i<CONFIG_KEYS.length; i++) {
			if (cfg.getProperty(CONFIG_KEYS[i]) == null) {
				log.fatal("The configuration file has not the key '" + CONFIG_KEYS[i] + "' set! Please edit the file '" + CONFIG_FILENAME + "' and add this configuration parameter. Exiting...");
				System.exit(1);				
			}
		}
		// SNP frequency parameter
		SNP_FREQ = Double.parseDouble(cfg.getProperty(VARIANT_SNP_FREQ_KEY)); 
		if (SNP_FREQ < 0 || SNP_FREQ > 1) {
			log.fatal("The configuration parameter "
					+ VARIANT_SNP_FREQ_KEY
					+ " should be in the range [0.0-1.0]! Please correct. Exiting...");
			System.exit(1);
		}
		if (SNP_FREQ > 0) {
			Variant v = new Snp('A');
			v.setFrequency(SNP_FREQ);
			VariantHandler.addVariant(v);
		}
		
		// INDEL frequency parameter
		INS_FREQ = Double.parseDouble(cfg.getProperty(VARIANT_INS_FREQ_KEY)); 
		if (INS_FREQ < 0 || INS_FREQ > 1) {
			log.fatal("The configuration parameter "
					+ VARIANT_INS_FREQ_KEY
					+ " should be in the range [0.0-1.0]! Please correct. Exiting...");
			System.exit(1);
		}
		if (INS_FREQ > 0) {
			Variant v = new Insertion('A', 0, 0);
			v.setFrequency(INS_FREQ);
			VariantHandler.addVariant(v);
		}
		DEL_FREQ = Double.parseDouble(cfg.getProperty(VARIANT_DEL_FREQ_KEY)); 
		if (DEL_FREQ < 0 || DEL_FREQ > 1) {
			log.fatal("The configuration parameter "
					+ VARIANT_DEL_FREQ_KEY
					+ " should be in the range [0.0-1.0]! Please correct. Exiting...");
			System.exit(1);
		}
		if (DEL_FREQ > 0) {
			Variant v = new Deletion(0, 0);
			v.setFrequency(DEL_FREQ);
			VariantHandler.addVariant(v);
		}

		// CNV frequency parameter
		CNV_III_FREQ = Double.parseDouble(cfg.getProperty(VARIANT_CNV_III_FREQ_KEY)); 
		if (CNV_III_FREQ < 0 || CNV_III_FREQ > 1) {
			log.fatal("The configuration parameter "
					+ VARIANT_CNV_III_FREQ_KEY
					+ " should be in the range [0.0-1.0]! Please correct. Exiting...");
			System.exit(1);
		}
		if (CNV_III_FREQ > 0) {
			Variant v = new Cnv();
			v.setFrequency(CNV_III_FREQ);
			VariantHandler.addVariant(v);
		}

		// Indel min/max length
		INDEL_MIN_LENGTH = Integer.parseInt(cfg.getProperty(INDEL_LENGTH_MIN_KEY));
		INDEL_MAX_LENGTH = Integer.parseInt(cfg.getProperty(INDEL_LENGTH_MAX_KEY));
		if (INDEL_MIN_LENGTH > INDEL_MAX_LENGTH) {
			log.fatal("The configuration parameter " + INDEL_LENGTH_MIN_KEY + " can't be greater than " + INDEL_LENGTH_MAX_KEY+"! Please correct. Exiting...");
			System.exit(1);
		}
		if (INDEL_MIN_LENGTH <= 0 || INDEL_MAX_LENGTH <= 0) {
			log.fatal("The configuration parameters " + INDEL_MIN_LENGTH + " and " + INDEL_MAX_LENGTH + " can't be zero or negative! Please correct. Exiting...");
			System.exit(1);
		}
		
		// FASTA file columns [1-120]
		FASTA_FILE_COLUMNS = Integer.parseInt(cfg.getProperty(FASTA_FILE_COLUMNS_KEY));
		if (FASTA_FILE_COLUMNS < 1 || FASTA_FILE_COLUMNS > 120) {
			log.fatal("The configuration parameter "
					+ FASTA_FILE_COLUMNS_KEY
					+ " should be in the range [1-120]! Please correct. Exiting...");
			System.exit(1);
		}
		
		// Load region sets
		int regionSetCount = Integer.parseInt(cfg.getProperty(REGION_SET_COUNT_KEY));
		if (regionSetCount > 0) {
			REGION_SETS = new RegionSet[regionSetCount];
			for (int i=1; i<=regionSetCount; i++) {
				String name = cfg.getProperty(REGION_SET_NAME_KEY.replaceAll("#", Integer.toString(i)));
				String freq = cfg.getProperty(REGION_SET_VARIANT_FREQ_KEY.replaceAll("#", Integer.toString(i)));
				String bed = cfg.getProperty(REGION_SET_BED_KEY.replaceAll("#", Integer.toString(i)));
				try {
					log.info("Loading and reordering BED file '" + bed + "'...");
					REGION_SETS[i-1] = new RegionSet(name, Double.parseDouble(freq), bed);
					
					// Keep a pointer to "reserved" name regions
					if (name.equals(REGION_SET_NAME_EXON))
						EXONS = REGION_SETS[i-1];
					else
					if (name.equals(REGION_SET_NAME_INTRON))
						INTRONS = REGION_SETS[i-1];
					else
					if (name.equals(REGION_SET_NAME_GENE))
						GENES = REGION_SETS[i-1];
					
				} catch (Exception e) {
					e.printStackTrace();
					log.fatal("An error occured while reading region set parameters from configuration file!\n" + e.getMessage() + "\nExiting...");
					System.exit(1);
				}
			}
		}
		DEFAULT_REGION_FREQ = Double.parseDouble(cfg.getProperty(REGION_SET_DEFAULT_FREQ_KEY)); 
		if (DEFAULT_REGION_FREQ < 0 || DEFAULT_REGION_FREQ > 1) {
			log.fatal("The configuration parameter "
					+ REGION_SET_DEFAULT_FREQ_KEY
					+ " should be in the range [0.0-1.0]! Please correct. Exiting...");
			System.exit(1);
		}
		
	}
	
	
	/**
	 * Constructor.
	 * Initialize.
	 * 
	 * @param inFasta the input FASTA file (reference genome)
	 * @param chr the chromosome name corresponding to the sequence to process
	 * @param outFasta the output FASTA file (reference genome with simulated variations)
	 * @param vcfFile the file to output variants in VCF format
	 */
	public DNAVSim(String inFasta, String chr, String outFasta, String vcfFile) {
		
		// Initialize logger
		log = Logger.getLogger("DNAVSim");
		log.setLevel(Level.DEBUG);
		
		log.info("Logger DNAVSim initialized");

		log.info("Reading reference genome from file '" + inFasta + "'...");
		log.info("Writinng modified genome to '" + outFasta + "'...");
		log.info("Writing variants to VCF file '" + vcfFile + "'...");
	
		loadConfiguration();

		// Counters for producing a report at the end of the process
		int snpCounter = 0;
		int indelCounter = 0;
		int cnvCounter = 0;
		int currentMbp = 0;
		
		String seqHeaderLine = "";
		try {
			seqHeaderLine = preformatFastaFile(inFasta, TMP_FILNAME_TAG);
		} catch (IOException e1) {
			log.fatal("An IO error occured while preformating FASTA file! Exiting...");
			System.exit(0);
		}
		
		// Open I/O files
		InputStream in = null;
		FastaFileWriter out = null;
		VCFWriter vcf = null;
		try {
			in = new BufferedInputStream(new FileInputStream(inFasta + TMP_FILNAME_TAG));
			out = new FastaFileWriter(new BufferedWriter(new FileWriter(
					outFasta)), Integer.parseInt(cfg
					.getProperty(FASTA_FILE_COLUMNS_KEY)));
			vcf = new VCFWriter(new BufferedWriter(new FileWriter(vcfFile)));
		} catch (FileNotFoundException e) {
			log.fatal("File not found: " + e.getMessage());
			System.exit(1);
		} catch (IOException e) {
			log.fatal("IO Exception when opening I/O files: " + e.getMessage());
			System.exit(1);
		}		
		
		log.info("Processing FASTA file...");
		
		// important variables: chr and pos gives the current locus of the AA
		// in a chromosome (absolute locus is therefore defined by <chr>:<pos>)
		// being read in inFasta file (original genome file)
//		String chr = parseFastaHeaderLine(seqHeaderLine);
		int pos = 0; 

		// Initialize random numbers generator
		Math.random();
		double r = Math.random();		

		int bufLen = 1024;
		byte[] buf = new byte[bufLen];

		try {

			out.printHeaderLine("> " + seqHeaderLine);
			
			int len = 0;
			while ((len = in.read(buf, 0, bufLen)) != -1) {
					
					// Respectively the REF and ALT column value to be written in VCF file
					char ref;
					String alt = "";
					
					for (int i=0; i<len; i++) {
						
						pos++;
						
						// Write current position if passed a Mbp
						if ((pos / 1000000) > currentMbp) {
							log.debug("visiting locus "+chr+":"+pos);
							currentMbp++;
						}
						
						// Determine the frequency of a variant at this position
						// checking the region sets defined by user
						double variantFrequency = -1;
						String info = "";
						for (RegionSet region:REGION_SETS) {
							if (region.hasFeature(chr, pos-1)) {
								variantFrequency = region.getVariantFrequency();
								if (info.length()>0) info += ";";
								info += "RegionSet=" + region.getName();
							}							
						}
						if (variantFrequency == -1) {
							variantFrequency = DEFAULT_REGION_FREQ;
							if (info.length()>0) info += ";";
							info += "RegionSet=DEFAULT";
						}						
						
						ref = (char) buf[i];

						if (ref != 'N') {

							r = Math.random();

							// A variant is to be created here
							if (r < variantFrequency) {

								// Choose a variant type to create
								Variant.TYPE vType = VariantHandler.chooseRandomVariant();								
									
								/**
								 * Insertion (INDEL)
								 */
								if (vType == Variant.TYPE.INSERTION) {
									Insertion ins = new Insertion(ref, INDEL_MIN_LENGTH, INDEL_MAX_LENGTH);
									alt =  ref + ins.getAlternate();
									out.printBases(alt);									
									vcf.printVariant(chr, Integer.toString(pos), Character.toString(ref), alt, info);
									indelCounter++;
								}
								else
									
								/**
								 * Deletion (INDEL)
								 */
								if (vType == Variant.TYPE.DELETION) {

									Deletion del = new Deletion(INDEL_MIN_LENGTH, INDEL_MAX_LENGTH);
									alt = Character.toString(ref);
									String refBases = Character.toString(ref);
									
									// Skip 'length' nucleotide(s) from input 
									for (int j=0; j<del.getLength(); j++) {
										i++;
										if (i == len) {
											len = in.read(buf, 0, bufLen);
											if (len == -1) break; // Reduce deletion length as we reached EOF
											i = 0;
										}
										refBases += (char)buf[i];
										
									}
									
									vcf.printVariant(chr, Integer.toString(pos), refBases, alt, info);
									out.printBase(ref);
									pos += del.getLength();
									indelCounter++;
								}
								else
								
								/**
								 * create CNV
								 */
								if (vType == Variant.TYPE.CNV) {
									
									// CNV type III
									//
									// For now the simulator test if it is in the middle of a gene 
									// test if there is an exon following and copy it n times
									BEDFeature gene = GENES.getBEDFeature(chr, pos-1);
									if (gene != null) {
										BEDFeature nextExon = EXONS.getNextBEDFeature(chr, pos-1);
										
										// Check if nextExon is still in the same gene
										if (nextExon != null && nextExon.getChromEnd() <= gene.getChromEnd()) {
											
											/**
											 * For now we copy 1000 times the current position up to the
											 * end of the next exon (experimental)
											 */
											int refLength = nextExon.getChromEnd() - pos;
											int copyFactor = 100;
											// Skip 'length' nucleotide(s) from input 
											String refBases = "";
											for (int j=0; j<refLength; j++) {
												i++;
												if (i == len) {
													len = in.read(buf, 0, bufLen);
													if (len == -1) break; // Reduce deletion length as we reached EOF
													i = 0;
												}
												refBases += (char)buf[i];
											}
											alt = "";
											for (int j=0; j<copyFactor; j++)
												alt += refBases;
											vcf.printVariant(chr, Integer.toString(pos), refBases, alt, info);
											out.printBases(refBases);
											pos += refLength;
											cnvCounter++;
										}
										
										// Abort with a 'no variation' event
										else {
											out.printBase(ref);
										}
										
									}
									
								}
								else
									
								/**
								 * create SNP
								 */
								if (vType == Variant.TYPE.SNP) {
									Snp s = new Snp((char)buf[i]);
									alt = s.getAlternate();
									vcf.printVariant(chr, Integer.toString(pos), Character.toString(ref), alt, info);
									out.printBases(alt);
									snpCounter++;
								}

								/**
								 * No variation
								 */
								else {
									out.printBase(ref);
								}
								
							}

							/**
							 * No variation
							 */
							else {
								out.printBase(ref);
							}

						}
						
						/**
						 * No variation
						 */
						else {
							out.printBase(ref);
						}
						
					}
					
			}
			
		} catch (IOException e) {
			log.fatal("IO Exception when reading genome FASTA file!\n" + e.getMessage());
			System.exit(1);
		}
		
		
		// Closing opened files
		try {
			in.close();
			out.close();
			vcf.close();
		} catch (IOException e) {
			log.fatal("IO Error when closing FASTA file!\n" + e.getMessage());
		}

		// Delete temporary files
		
		
		// Print report  
		log.info("Report:");
		log.info("-------");
		log.info(" SNP(s) created:   " + snpCounter);
		log.info(" Indel(s) created: " + indelCounter);
		log.info(" CNV(s) created:   " + cnvCounter);
		log.info("Done. Exiting...");

	}
	
	/**
	 * @param args
	 */
	public static void main(String[] args) {

		
		if (args.length < 4) {
			System.out.println("Usage: java DNAVSim [genome FASTA file name] [chromosome name (e.g. 1,2,...,X,Y,MT] [output file (FASTA file)] [output VCF file]");
			System.out.println("");
			System.out.println("Edit the configuration file 'dnavsim.cfg' to adapt the parameters (regions, frequencies, etc.).");
			System.exit(1);
		}
		DNAVSim sim = new DNAVSim(args[0], args[1], args[2], args[3]);
		
	}

	
	/**
	 * Read a single sequence FASTA file, remove header line (e.g. ">") and all newline character
	 * then write the result to the file name with a given suffix.
	 * 
	 * @param fasta the FASTA file name to process
	 * @param tmpSuffix the temporary suffix to be appended to the file name 
	 * @return the header line content
	 * @throws IOException
	 */
	private String preformatFastaFile(String fasta, String tmpSuffix) throws IOException {

		log.info("Reformating FASTA file '" + fasta + "' removing newline characters and sequence header line. Result to be outputed in file '" + fasta + tmpSuffix + "'...");
		boolean seqHdrRead = false;
		InputStream in = null;
		OutputStream out = null;
		String headerLine = "";

		in = new BufferedInputStream(new FileInputStream(fasta));
		out = new BufferedOutputStream(new FileOutputStream(fasta + tmpSuffix));
		int bufLen = 10 * 1024;
		byte[] buf = new byte[bufLen];
		int len = 0;
		while ((len = in.read(buf, 0, bufLen)) != -1) {
			int i = 0;
			int j = 0;
			for (j = 0; j < len; j++) {

				// Sequence sequence header lines
				if (buf[j] == '>') {
					if (seqHdrRead)
						throw new IOException(
								"The FASTA file contains more than 1 sequence! Sorry but this is not supported yet.");
					seqHdrRead = true;
					int hdrStart = j + 1;
					while (buf[j] != 0x0A)
						j++;
					int hdrEnd = j;
					headerLine = new String(buf, hdrStart, hdrEnd - hdrStart);
					log.info("Reading sequence '" + headerLine + "'...");
					j++;
					i = j;
				}

				if (buf[j] == 0x0A) {
					out.write(buf, i, (j - i));
					i = j + 1;
				}

			}
			if (i < j + 2) {
				out.write(buf, i, (j - i));
			}
		}
		in.close();
		out.close();

		log.info("Reformating done.");
		return headerLine;
	}
	
}
